#include <src/pcps.h>
#include <src/input.h>
#include <src/correlator.h>
#include <src/wavefunction.h>
#include <src/operator.h>
#include <src/hamiltonian.h>
#include <src/timing.h>
#include <src/projected_optimization_ti.h>

//--------------------------------------------------------------------------------
// pcps::solve_projected_nr_equations -- Solves the newton raphson equations
//                                       of projected CPS for the update vector,
//                                       which is returned in the residual.
//                                       The jacobian data may be destroyed.
//--------------------------------------------------------------------------------

template <class SCALAR>
void pcps::solve_projected_nr_equations(const pcps::Input & userinp, const int n, SCALAR * jacobian, SCALAR * residual) {

      // get the dimension of the real jacobian
      const int N = ( typeid(SCALAR) == typeid(double) ? n : 2*n );

      boost::shared_array<double> jac_real;
      boost::shared_array<double> res_real;

      double * jac = (double *)jacobian;
      double * res = (double *)residual;

      // convert the complex jacobian and residual to real form
      if ( typeid(SCALAR) == typeid(std::complex<double>) ) {

        // allocate the real jacobian and residual
        jac_real = pcps::allocate_shared_array<double>(N*N);
        res_real = pcps::allocate_shared_array<double>(N);
        jac = jac_real.get();
        res = res_real.get();

        // convert residual
        for (int i = 0; i < n; i++) {
          res[i  ] = pcps::real(residual[i]);
          res[i+n] = pcps::imag(residual[i]);
        }

        // convert jacobian
        for (int i = 0; i < n; i++)
        for (int j = 0; j < n; j++) {
          jac[ (i+0)*N + j+0 ] =  pcps::real(jacobian[i*n+j]);
          jac[ (i+0)*N + j+n ] = -pcps::imag(jacobian[i*n+j]);
          jac[ (i+n)*N + j+0 ] =  pcps::imag(jacobian[i*n+j]);
          jac[ (i+n)*N + j+n ] =  pcps::real(jacobian[i*n+j]);
        }

      }

      // get the SVD of the jacobian
      const int lwork = 5*N;
      boost::shared_array<double> u     = pcps::allocate_shared_array<double>(N*N);
      boost::shared_array<double> vt    = pcps::allocate_shared_array<double>(N*N);
      boost::shared_array<double> work  = pcps::allocate_shared_array<double>(lwork);
      boost::shared_array<double> s     = pcps::allocate_shared_array<double>(N);
      boost::shared_array<double> rwork = pcps::allocate_shared_array<double>(5*N);
      {
        int info;
        pcps::xgesvd('A', 'A', N, N, jac, N, s.get(), vt.get(), N, u.get(), N, work.get(), lwork, rwork.get(), info);
        if (info != 0)
          throw pcps::Exception( (boost::format("xgesvd failed in pcps::solve_projected_nr_equations with error code %i") % info).str() );
        //// check that the jacobian was computed correctly (jacobian = u * s * vt)
        //boost::shared_array<SCALAR> new_jac = pcps::allocate_shared_array<SCALAR>(N*N, pcps::zero<SCALAR>());
        //for (int i = 0; i < N; i++)
        //for (int j = 0; j < N; j++)
        //for (int k = 0; k < N; k++)
        //  new_jac[i*N+j] += u[i*N+k] * s[k] * vt[k*N+j];
        //for (int i = 0; i < N*N; i++)
        //  new_jac[i] -= saved_jac[i];
        //double max_error = 0.0;
        //for (int i = 0; i < N*N; i++)
        //  if (std::abs(new_jac[i]) > max_error)
        //    max_error = std::abs(new_jac[i]);
        //if (userinp.verbose_print())
        //  std::cout << boost::format("max error in new_jac = %.4e") % max_error << std::endl;
      }

      // print the jacobian's singular values
      if (userinp.verbose_print()) {
        std::cout << "jacobian singular values:" << std::endl;
        for (int i = 0; i < N; i++)
          std::cout << boost::format("%6i  %s") % i % pcps::format_number("%18.6e", s[i]) << std::endl;
        std::cout << std::endl;
      }

      // compute the update vector and place it in the residual
      const double shift = 0.0;
      int m = 0;
      for (int i = 0; i < N && s[i] > 1.0e-10; i++, m++) {}
      pcps::xgemm('N', 'N', m, 1, N, 1.0, u.get(), N, res, N, 0.0, work.get(), m);
      for (int i = 0; i < m; i++)
        work[i] = -work[i] / ( s[i] + shift );
      pcps::xgemm('N', 'N', N, 1, m, 1.0, vt.get(), N, work.get(), m, 0.0, res, N);

      // convert back to complex form
      if ( typeid(SCALAR) == typeid(std::complex<double>) )
        for (int i = 0; i < n; i++)
          residual[i] = pcps::unity<SCALAR>() * res[i] + pcps::imaginary_unity<SCALAR>() * res[i+n];

}

//--------------------------------------------------------------------------------
// pcps::projected_driver_ti -- performs a projected cps optimization of the
//                              wavefunction using translational invariance
//--------------------------------------------------------------------------------

template <class SCALAR>
void pcps::projected_driver_ti(const pcps::Input & userinp, const pcps::Hamiltonian<SCALAR> & ham, pcps::Wavefunction<SCALAR> & wfn) {

  // get MPI info
  const MPI::Comm & comm = MPI::COMM_WORLD;
  const int nproc = comm.Get_size();
  const int myrank = comm.Get_rank();

  // set limits of macro iteration loop
  const int maxloop = userinp.max_macro_loop();
  const double tol = userinp.energy_tolerance();

  // start a timer
  pcps::Stopwatch sw; sw.start(); 

  // print a greeting
  if (myrank == 0) {
    std::cout << "  --- Performing a Translationally Invariant Projective Optimization of the Wavefunction ---" << std::endl;
    std::cout << std::endl;
  }

  // check that the particles are spin-1/2 spins
  if (userinp.particle_type() != pcps::keyword::Part_Spin)
    throw pcps::Exception( (boost::format("pcps::projected_driver_ti requires particle_type to be spin")).str() );
  if (userinp.configs_per_site() != 2)
    throw pcps::Exception( (boost::format("pcps::projected_driver_ti requires spins to be 1/2")).str() );

  //// check that the lattice is cartesian
  //if (userinp.lattice_type() != pcps::keyword::Lattice_Cartesian)
  //  throw pcps::Exception( (boost::format("pcps::projected_driver_ti requires the lattice to be cartesian")).str() );

  // check that the lattice is two dimensional
  if (userinp.lattice_length_z() != 1)
    throw pcps::Exception( (boost::format("pcps::projected_driver_ti requires lattice_length_z to be 1")).str() );

  // check that the boundary is periodic
  if (userinp.boundary() != pcps::keyword::Boundary_Periodic)
    throw pcps::Exception( (boost::format("pcps::projected_driver_ti requires periodic boundary conditions")).str() );

  // get the lattice's x and y dimensions
  const int lx = userinp.lattice_length_x();
  const int ly = userinp.lattice_length_x();

  // get the unique Hamiltonian operators
  std::vector<typename pcps::OpBase<SCALAR>::const_ptr > ops;
  if (userinp.lattice_type() == pcps::keyword::Lattice_Cartesian) {

    ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpRL      <SCALAR>( 0,  1,  0.50 * pcps::unity<SCALAR>()) ) );
    ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpRL      <SCALAR>( 1,  0,  0.50 * pcps::unity<SCALAR>()) ) );
    ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpNN      <SCALAR>( 0,  1,  1.00 * pcps::unity<SCALAR>()) ) );
    ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpN       <SCALAR>(     0, -0.50 * pcps::unity<SCALAR>()) ) );
    ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpN       <SCALAR>(     1, -0.50 * pcps::unity<SCALAR>()) ) );
    ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpIdentity<SCALAR>(         0.25 * pcps::unity<SCALAR>()) ) );

    ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpRL      <SCALAR>( 0, lx,  0.50 * pcps::unity<SCALAR>()) ) );
    ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpRL      <SCALAR>(lx,  0,  0.50 * pcps::unity<SCALAR>()) ) );
    ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpNN      <SCALAR>( 0, lx,  1.00 * pcps::unity<SCALAR>()) ) );
    ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpN       <SCALAR>(     0, -0.50 * pcps::unity<SCALAR>()) ) );
    ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpN       <SCALAR>(    lx, -0.50 * pcps::unity<SCALAR>()) ) );
    ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpIdentity<SCALAR>(         0.25 * pcps::unity<SCALAR>()) ) );

  } else if (userinp.lattice_type() == pcps::keyword::Lattice_Triangular) {

    int a, b;
    for (int i = 0; i < 3; i++) {

      a = 0+i;
      b = 1+i;
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpRL      <SCALAR>( a, b,  0.50 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpRL      <SCALAR>( b, a,  0.50 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpNN      <SCALAR>( a, b,  1.00 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpN       <SCALAR>(    a, -0.50 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpN       <SCALAR>(    b, -0.50 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpIdentity<SCALAR>(        0.25 * pcps::unity<SCALAR>()) ) );

      a = 0+i;
      b = lx+i;
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpRL      <SCALAR>( a, b,  0.50 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpRL      <SCALAR>( b, a,  0.50 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpNN      <SCALAR>( a, b,  1.00 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpN       <SCALAR>(    a, -0.50 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpN       <SCALAR>(    b, -0.50 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpIdentity<SCALAR>(        0.25 * pcps::unity<SCALAR>()) ) );

      a = 0+i;
      b = lx+1+i;
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpRL      <SCALAR>( a, b,  0.50 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpRL      <SCALAR>( b, a,  0.50 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpNN      <SCALAR>( a, b,  1.00 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpN       <SCALAR>(    a, -0.50 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpN       <SCALAR>(    b, -0.50 * pcps::unity<SCALAR>()) ) );
      ops.push_back( typename pcps::OpBase<SCALAR>::const_ptr( new pcps::OpIdentity<SCALAR>(        0.25 * pcps::unity<SCALAR>()) ) );

    }

  } else throw pcps::Exception( (boost::format("pcps::projected_driver_ti requires the lattice to be cartesian or tirangular")).str() );

  //// find the variable correlators in the wavefunction
  //std::vector<pcps::Correlator<SCALAR> > var_corrs;
  //for (typename std::vector<pcps::Correlator<SCALAR> >::iterator corr = wfn.correlators().begin(); corr != wfn.correlators().end(); corr++)
  //  if (!corr->fixed() && !corr->ti())
  //    var_corrs.push_back(*corr);
  //  else if (!corr->fixed() && corr->ti())
  //    throw pcps::Exception( (boost::format("pcps::projected_driver_ti requires all variable correlators to not be ti images")).str() );

  // ensure no wavefunction correlators are maked as ti (all non-fixed correlators are treated as being translationally invariant)
  for (typename std::vector<pcps::Correlator<SCALAR> >::iterator corr = wfn.correlators().begin(); corr != wfn.correlators().end(); corr++)
    if (!corr->fixed() && corr->ti())
      throw pcps::Exception( (boost::format("pcps::projected_driver_ti requires all variable correlators to not be ti images")).str() );

  // for each unique operator, build the necessary cluster infrastructure
  std::vector<boost::shared_ptr<pcps::TiOpCluster<SCALAR> > > op_clusters;
  for (typename std::vector<typename pcps::OpBase<SCALAR>::const_ptr>::const_iterator op = ops.begin(); op != ops.end(); op++)
    op_clusters.push_back( boost::shared_ptr<pcps::TiOpCluster<SCALAR> >( new pcps::TiOpCluster<SCALAR>(userinp, lx, ly, *op, wfn) ) );

  // get the number of variables
  const int nvar = wfn.nvars();

  // initialize energy, energy derivatives, residual, and jacobian
  SCALAR energy = pcps::zero<SCALAR>();
  std::vector<SCALAR> energy_deriv(nvar, pcps::zero<SCALAR>());
  std::vector<SCALAR> residual(nvar, pcps::zero<SCALAR>());
  std::vector<SCALAR> jacobian(nvar*nvar, pcps::zero<SCALAR>());

  // initialize a vector for reducing
  std::vector<SCALAR> red_target( ( myrank == 0 ? nvar*nvar : 1 ), pcps::zero<SCALAR>());

  // perform iterative newton raphson optimization
  bool is_converged = false;
  for (int iter = 0; iter < maxloop && !is_converged; iter++) {

    // start a timer for the macro iteration
    pcps::Stopwatch sw_mac; sw_mac.start(); 

    // record the energy of the previous iteration
    const SCALAR last_energy = energy;

    // zero the energy, derivatives, residual, and jacobian
    energy = pcps::zero<SCALAR>();
    pcps::xscal(energy_deriv.size(), pcps::zero<SCALAR>(), &energy_deriv.at(0), 1);
    pcps::xscal(residual.size(), pcps::zero<SCALAR>(), &residual.at(0), 1);
    pcps::xscal(jacobian.size(), pcps::zero<SCALAR>(), &jacobian.at(0), 1);

    // for each unique operator, loop over all configurations of that operator's cluster
    for (int i = 0; i < op_clusters.size(); i++) {

      op_clusters.at(i)->reset();

      op_clusters.at(i)->compute_all_contributions(userinp);

      op_clusters.at(i)->add_contributions(userinp, ops.at(i)->coeff(), energy, energy_deriv, residual, jacobian);
      
    }

    // reduce energy
    pcps::reduce(comm, &energy, &red_target.at(0), 1, MPI::SUM, 0);
    if (myrank == 0) energy = red_target.at(0);

    // reduce energy derivatives
    pcps::reduce(comm, &energy_deriv.at(0), &red_target.at(0), energy_deriv.size(), MPI::SUM, 0);
    if (myrank == 0) pcps::xcopy(energy_deriv.size(), &red_target.at(0), 1, &energy_deriv.at(0), 1);

    // reduce residual
    pcps::reduce(comm, &residual.at(0), &red_target.at(0), residual.size(), MPI::SUM, 0);
    if (myrank == 0) pcps::xcopy(residual.size(), &red_target.at(0), 1, &residual.at(0), 1);

    // reduce jacobian
    pcps::reduce(comm, &jacobian.at(0), &red_target.at(0), jacobian.size(), MPI::SUM, 0);
    if (myrank == 0) pcps::xcopy(jacobian.size(), &red_target.at(0), 1, &jacobian.at(0), 1);

    // on root node, update the correlators
    if (myrank == 0) {

      // loop over projectors' occupations to get overlap elements
      for (int i = 0; i < wfn.correlators().size(); i++) {

        // skip fixed correlators
        if (wfn.correlators().at(i).fixed()) continue;

        // create a cluster and density matrix for the projector
        pcps::Cluster pocc(userinp.nsites(), userinp.configs_per_site(), wfn.correlators().at(i).site_vec());
        pcps::Compact_RDM<SCALAR, pcps::OpRL<SCALAR>, pcps::Part_Spin, pcps::Ref_Uniform> rdm(userinp, wfn, pcps::OpRL<SCALAR>(), pocc);

        // loop over the projector's configurations
        for ( ; !pocc.finished(); pocc++) {

          // get overlap matrix element
          const SCALAR ome = rdm.element(pocc);

          // get compound index
          const int cmpd = wfn.corr_offsets().at(i) + wfn.correlators().at(i).get_cmpd(pocc);

          // update residual
          residual[cmpd] -= energy * ome;

          // update jacobian
          pcps::xaxpy(nvar, -ome, &energy_deriv.at(0), 1, &jacobian.at(cmpd*nvar), 1);

        }

      }

      // print energy and residual norm
      std::cout << boost::format("macro iter %3i  energy = %s   residual = %12.4e   elapsed time = %17.6f")
                   % iter
                   % pcps::format_number("%20.12f", energy)
                   % std::abs(std::sqrt(pcps::xdotc(residual.size(), &residual.at(0), 1, &residual.at(0), 1)))
                   % pcps::get_overall_timer().elapsed_seconds()
                << std::endl;
      std::cout << std::endl;

      // print residual
      if (userinp.verbose_print()) {
        std::cout << "residual:" << std::endl;
        for (int i = 0; i < nvar; i++)
          std::cout << boost::format("%6i  %s") % i % pcps::format_number("%18.6f", residual[i]) << std::endl;
        std::cout << std::endl;
      }

      // solve the newton raphson equation for the update (which overwrites the residual)
      pcps::solve_projected_nr_equations(userinp, nvar, &jacobian.at(0), &residual.at(0));

      // add the update to the wavefunction and normalize the correlators
      for (int i = 0; i < wfn.correlators().size(); i++) {

        // skip fixed correlators
        if (wfn.correlators().at(i).fixed()) continue;

        // add the update to the correlator
        pcps::xaxpy(wfn.correlators().at(i).nelements(), pcps::unity<SCALAR>() * ( 1.0 - userinp.damping() ),
                    &residual.at(wfn.corr_offsets().at(i)), 1, wfn.correlators().at(i).data_ptr(), 1);

        // normalize the correlator
        wfn.correlators().at(i).normalize();

        // print the updated correlator
        if (userinp.verbose_print()) {
          pcps::print_correlator(wfn.correlators().at(i));
          std::cout << std::endl;
        }

      }

      // save the wavefunction to disk
      wfn.dump_binary(userinp);

    }

    // broadcast new variables to all nodes
    for (int i = 0; i < wfn.correlators().size(); i++)
      if (!wfn.correlators().at(i).fixed())
        comm.Bcast(wfn.correlators().at(i).data_ptr(), wfn.correlators().at(i).nelements() * sizeof(SCALAR), MPI::CHAR, 0);

    // check convergence
    is_converged = ( std::abs(energy - last_energy) < tol );
    comm.Bcast(&is_converged, 1, MPI::BOOL, 0);

  }

  // print the result of the iteration
  if (myrank == 0) {
    if (is_converged)
      std::cout << std::endl
                << boost::format("pcps macro iterations converged!") << std::endl
                << std::endl;
    else
      std::cout << std::endl
                << boost::format("no convergence...") << std::endl
                << std::endl;
  }

  // print the elapsed time
  if (myrank == 0) {
    std::cout << boost::format("time to optimize wavefunction: %.2f sec") % sw.elapsed_seconds() << std::endl;
    std::cout << std::endl;
  }

}

// explicitly instantiate the template
template void pcps::projected_driver_ti(const pcps::Input &,
                                        const pcps::Hamiltonian<double> &,
                                        pcps::Wavefunction<double> &);
template void pcps::projected_driver_ti(const pcps::Input &,
                                        const pcps::Hamiltonian< std::complex<double> > &,
                                        pcps::Wavefunction< std::complex<double> > &);
